import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import FancyBboxPatch, Rectangle, ConnectionPatch


def create_professional_architecture():
    plt.rcParams.update({
        'font.family': 'DejaVu Sans',
        'mathtext.fontset': 'dejavusans',
        'axes.titlesize': 14,
        'axes.labelsize': 12,
        'font.size': 11,
        'figure.dpi': 300
    })

    # 扩大画布尺寸，为模块提供更多空间
    fig, ax = plt.subplots(figsize=(14, 10))
    ax.set_xlim(0, 120)
    ax.set_ylim(0, 80)
    ax.axis('off')

    COLORS = {
        'backbone': {'fc': '#F5F5F5', 'ec': '#4B4B4B', 'lw': 2},
        'original': {'fc': '#E3F2FD', 'ec': '#1976D2', 'lw': 1.5},
        'improved': {'fc': '#E8F5E9', 'ec': '#2E7D32', 'lw': 1.8},
        'optim': {'fc': '#FFF8E1', 'ec': '#FF8F00', 'lw': 1.8},
        'highlight': {'fc': '#FFEBEE', 'ec': '#C62828', 'lw': 2}
    }

    def create_module(x, y, width, height, text, style, subtext=None):
        box = FancyBboxPatch(
            (x, y), width, height,
            boxstyle=f"round,pad=0.4,rounding_size=1.2",
            **style
        )
        ax.add_patch(box)

        lines = text.split('\n')
        if len(lines) > 1:
            ax.text(x + width / 2, y + height - 1.5, lines[0],
                    ha='center', va='top',
                    fontsize=10, fontweight='bold')
            ax.text(x + width / 2, y + height - 3.5, lines[1],
                    ha='center', va='top',
                    fontsize=9)
        else:
            ax.text(x + width / 2, y + height - 2, text,
                    ha='center', va='top',
                    fontsize=10, fontweight='bold')

        if subtext:
            ax.text(x + width / 2, y + height / 2 - 1, subtext,
                    ha='center', va='center',
                    fontsize=9)

    # 调整主干网络位置，增加垂直间距
    create_module(10, 60, 20, 10, "Input Image\n(512×512×3)", COLORS['backbone'],
                  r"$\mathbf{I} \in \mathbb{R}^{H×W×3}$")
    create_module(10, 35, 20, 10, "Backbone\n(ResNet-50)", COLORS['backbone'],
                  r"$\mathbf{F} \in \mathbb{R}^{\frac{H}{4}×\frac{W}{4}×256}$")
    create_module(10, 10, 20, 10, "Detection Head", COLORS['original'],
                  r"$\mathbf{P} = \Phi(\mathbf{F}^*)$")

    # 调整改进模块位置，增加水平和垂直间距
    create_module(45, 50, 22, 12, "LAM\n(Local Affinity)", COLORS['improved'],
                  r"$\alpha_{ij} = \frac{e^{q_i^T k_j}}{\sum e^{q_i^T k_j}}$")

    # LAM子模块改为两行排列
    for i, (dx, dy, text) in enumerate(zip([0, 8, 16], [0, 0, -8], ["Channel", "Spatial", "Fusion"])):
        create_module(47 + dx, 52 + dy, 6, 6, text, COLORS['highlight'],
                      r"$\alpha_{c/s}$" if i < 2 else r"$\otimes$")

    create_module(45, 25, 22, 12, "LCM\n(Local Collaborating)", COLORS['improved'],
                  r"$\beta_{ij} = \mathrm{softmax}(\frac{QK^T}{\sqrt{d}})$")

    # LCM子模块改为两行排列
    for i, (dx, dy, text) in enumerate(zip([0, 8, 16], [0, 0, -8], ["Reduction", "Self-Att", "Cross-Int"])):
        create_module(47 + dx, 27 + dy, 6, 6, text, COLORS['highlight'],
                      r"$\downarrow$" if i == 0 else r"$\beta$")

    # ADEM模块位置调整
    create_module(80, 35, 22, 12, "ADEM\n(Density Estimator)", COLORS['optim'],
                  r"$\mathcal{L}_{dens} = \|\rho - \hat{\rho}\|_2^2$")

    # ADEM子模块改为两行排列
    for i, (dx, dy, text) in enumerate(zip([0, 8, 16], [0, 0, -8], ["Density", "Uncertainty", "Weighting"])):
        create_module(82 + dx, 37 + dy, 6, 6, text, COLORS['optim'],
                      r"$\sigma^2$" if i == 1 else r"$\omega$")

    def connect_modules(start, end, style, text=None):
        con = ConnectionPatch(
            start, end,
            coordsA="data", coordsB="data",
            arrowstyle=f"->,head_width={style['head'] * 1.2},head_length={style['head'] * 2}",
            lw=style['lw'],
            linestyle=style.get('ls', '-'),
            color=style['color']
        )
        ax.add_artist(con)
        if text:
            mid = ((start[0] + end[0]) / 2, (start[1] + end[1]) / 2)
            ax.text(*mid, text, ha='center', va='center',
                    fontsize=10, backgroundcolor='white')

    # 调整连接线位置
    connect_modules((20, 60), (20, 45), {'color': '#4B4B4B', 'lw': 2.5, 'head': 5})
    connect_modules((20, 35), (20, 20), {'color': '#4B4B4B', 'lw': 2.5, 'head': 5})
    connect_modules((30, 50), (45, 50), {'color': '#2E7D32', 'lw': 1.8, 'head': 4},
                    r"$\mathbf{F}_a$")
    connect_modules((30, 25), (45, 25), {'color': '#2E7D32', 'lw': 1.8, 'head': 4},
                    r"$\mathbf{F}_b$")
    connect_modules((67, 45), (80, 35), {'color': '#FF8F00', 'lw': 1.8, 'head': 4, 'ls': '--'},
                    r"$\nabla \mathcal{L}$")

    legend_elements = [
        FancyBboxPatch((0, 0), 1, 1, fc=COLORS['backbone']['fc'], ec=COLORS['backbone']['ec'],
                       label='Backbone Components'),
        FancyBboxPatch((0, 0), 1, 1, fc=COLORS['original']['fc'], ec=COLORS['original']['ec'],
                       label='Original Modules'),
        FancyBboxPatch((0, 0), 1, 1, fc=COLORS['improved']['fc'], ec=COLORS['improved']['ec'],
                       label='Proposed Modules'),
        plt.Line2D([0], [0], color='#4B4B4B', lw=2.5, label='Feature Flow'),
        plt.Line2D([0], [0], color='#2E7D32', lw=1.8, label='Attention Flow'),
        plt.Line2D([0], [0], color='#FF8F00', lw=1.8, ls='--', label='Gradient Flow')
    ]

    ax.legend(handles=legend_elements,
              loc='lower right',
              framealpha=1,
              edgecolor='gray',
              title='Legend:',
              title_fontsize=12,
              fontsize=10,
              handlelength=1.5)

    plt.title("Improved DiffusionDet Architecture with Local Attention Mechanisms",
              fontsize=14, pad=20)

    plt.savefig('professional_architecture.pdf', bbox_inches='tight')
    plt.savefig('professional_architecture.png', dpi=600, bbox_inches='tight')
    plt.show()


create_professional_architecture()
